import numpy as np
from ultralytics import YOLO
import os
# 抑制libpng警告
os.environ['OPENCV_LOG_LEVEL'] = 'ERROR'
import cv2
from PIL import Image, ImageDraw, ImageFont
import time
import copy


class VehicleDetector:
    def __init__(self, model_path="./models/yolov8m.pt", device="0"):
        """
        初始化车辆检测器
        :param model_path: YOLO模型路径
        :param device: 运行设备，"0"表示GPU，"cpu"表示CPU
        """
        # 加载YOLOv8模型，并设置verbose=False关闭输出
        self.model = YOLO(model_path, verbose=False)
        self.device = device

        # 定义车辆相关类别 (YOLOv8的COCO数据集类别)
        self.vehicle_classes = {
            2: "car",  # 小汽车
            3: "motorcycle",  # 摩托车
            5: "bus",  # 公交车
            7: "truck"  # 卡车
        }

        # 定义不同类别对应的颜色（使用中文键名）
        self.colors = {
            "car": (0, 255, 0),  # 绿色
            "motorcycle": (0, 0, 255),  # 红色
            "bus": (255, 255, 255),  # 白色
            "truck": (255, 255, 0)  # 青色
        }
        
        # 设置中文字体
        self.font_path = self._get_valid_font_path()
        self.font_size = 24  # 增大默认字体大小
        self.bg_color = (0, 0, 0, 180)  # 默认背景色：半透明黑色 (R, G, B, Alpha)
        
        # FPS计算相关变量
        self.prev_time = 0
        self.fps = 0
        self.fps_count = 0
        self.fps_start_time = time.time()
        self.fps_update_interval = 1.0  # 每秒更新一次FPS
        self.fps_display_interval = 5  # 每5帧更新一次FPS显示，减少转换频率
        self.frame_count = 0

    def _get_valid_font_path(self):
        """
        获取有效的中文字体路径
        :return: 字体文件路径
        """
        # 尝试多个可能的字体路径
        possible_paths = [
            "./fonts/simhei.ttf",  # 用户指定的路径
            "C:/Windows/Fonts/simhei.ttf",  # Windows系统黑体
            "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",  # Linux系统字体
            "/System/Library/Fonts/PingFang.ttc"  # macOS系统字体
        ]
        
        for path in possible_paths:
            if os.path.exists(path):
                print(f"找到字体文件: {path}")
                return path
        
        print("警告: 未找到中文字体文件，将使用默认字体")
        return None

         
    def batch_add_chinese_text(self, img, texts_info):
        """
        批量在OpenCV图像上添加中文文本，提高多文本绘制时的性能
        :param img: OpenCV图像
        :param texts_info: 文本信息列表，每个元素为(text, position, color, font_size, bg_color)的元组
        :return: 添加所有文本后的图像
        """
        # 验证输入图像
        if img is None or img.size == 0:
            print("警告: 输入图像为空")
            return img
            
        try:
            # 只进行一次图像格式转换
            img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(img_pil, "RGBA")
            
            # 遍历所有需要添加的文本信息
            for text_info in texts_info:
                text, position, color, font_size, bg_color = text_info
                
                # 使用默认背景色或传入的背景色
                if bg_color is None:
                    bg_color = self.bg_color
                
                font = ImageFont.truetype(self.font_path, font_size)
                    
                # 计算文本尺寸 - 兼容新旧版本的Pillow
                
                # 尝试使用新版的textbbox方法（Pillow 8.0+）
                bbox = draw.textbbox((0, 0), text, font=font)
                text_width = bbox[2] - bbox[0]
                text_height = bbox[3] - bbox[1]

                
                # 计算背景框位置和大小，并确保在图像范围内
                x, y = position
                img_height, img_width = img.shape[:2]
                
                # 确保背景框不会超出图像边界，并确保bg_y1 < bg_y2
                bg_x1 = max(0, x)
                bg_y1 = max(0, y - text_height)
                bg_x2 = min(img_width, x + text_width + 8)  # 右边距+8
                bg_y2 = min(img_height, y + 4)  # 下边距+4
                
                # 确保y1始终小于y2，避免矩形绘制错误
                if bg_y1 >= bg_y2:
                    # 如果文本位置太靠近顶部，将文本放在检测框下方
                    bg_y1 = y
                    bg_y2 = min(img_height, y + text_height + 4)
                
                # 绘制半透明背景
                bg_rgba = (bg_color[2], bg_color[1], bg_color[0], bg_color[3])  # 转换为RGBA
                draw.rectangle([(bg_x1, bg_y1), (bg_x2, bg_y2)], fill=bg_rgba)
                
                # 绘制文本 - 修复文本位置计算，确保文本在背景框内
                text_x = max(0, x + 4)
                # 计算文本的Y坐标，确保文本位于背景框内
                # 对于PIL的text方法，Y坐标是文本基线的位置，需要根据字体特性调整
                if bg_y1 < y - text_height:
                    # 文本在正常位置（检测框上方）
                    text_y = bg_y1 + text_height
                else:
                    # 文本在检测框下方
                    text_y = bg_y1 + text_height
                
                # 确保文本Y坐标不会超出图像范围
                text_y = min(img_height, text_y)
                
                # 绘制文本
                draw.text((text_x, text_y), text, font=font, fill=(color[2], color[1], color[0]))  # PIL使用RGB颜色
            
            # 只进行一次图像格式转换回OpenCV
            result_img = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
            return result_img
            
        except Exception as e:
            print(f"批量添加中文文本时出错: {str(e)}")
            # 出错时返回原始图像
            return img

    def _calculate_fps(self):
        """
        计算并更新FPS值
        """
        current_time = time.time()
        self.fps_count += 1
        self.frame_count += 1
        
        # 计算帧率（每帧的时间差）
        if self.prev_time > 0:
            frame_time = current_time - self.prev_time
            if frame_time > 0:
                current_fps = 1.0 / frame_time
                # 使用移动平均来平滑FPS值
                self.fps = 0.9 * self.fps + 0.1 * current_fps
        
        self.prev_time = current_time
        
        # 定期更新FPS显示值
        if current_time - self.fps_start_time >= self.fps_update_interval:
            self.fps_start_time = current_time

    def detect_and_draw(self, image_path, conf_threshold=0.3, show_fps=True):
        """
        检测图片文件中的车辆并绘制边界框（原文件路径输入模式）
        :param image_path: 输入图片路径
        :param conf_threshold: 置信度阈值
        :param show_fps: 是否显示FPS信息
        :return: 标绘后的图像数组，如果出错则返回None
        """
        try:
            # 检查图片文件是否存在
            if not os.path.exists(image_path):
                print(f"错误: 图片文件不存在 - {image_path}")
                return None

            # 读取图片
            image = cv2.imread(image_path)
            if image is None:
                print(f"错误: 无法读取图片 - {image_path}")
                return None

            # 调用帧处理方法
            return self.detect_frame_and_draw(image, conf_threshold, show_fps)

        except Exception as e:
            print(f"检测过程中发生错误: {str(e)}")
            return None

    def detect_frame_and_draw(self, frame, conf_threshold=0.3, show_fps=True):
        """
        检测图像帧中的车辆并绘制边界框（RTSP解码帧输入模式）
        :param frame: 输入的图像帧（numpy数组）
        :param conf_threshold: 置信度阈值
        :param show_fps: 是否显示FPS信息
        :return: 标绘后的图像数组，如果出错则返回None
        """
        try:
            # 检查输入是否为有效的numpy数组
            if not isinstance(frame, np.ndarray):
                print("错误: 输入不是有效的numpy数组")
                return None

            # 验证图像尺寸和类型
            if frame.size == 0 or len(frame.shape) < 2:
                print("错误: 输入图像无效")
                return None

            # 确保图像是3通道BGR格式
            if len(frame.shape) == 2:
                frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
            elif frame.shape[2] == 4:
                frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)

            # 执行检测，添加verbose=False关闭输出
            results = self.model(frame, device=self.device, conf=conf_threshold, verbose=False)

            # 收集所有需要绘制的文本信息，用于批量处理
            texts_info = []

            # 处理检测结果
            for result in results:
                # 获取检测框信息
                boxes = result.boxes

                # 遍历所有检测框
                for box in boxes:
                    # 获取类别ID和置信度
                    cls_id = int(box.cls[0])
                    confidence = float(box.conf[0])

                    # 只处理车辆相关类别
                    if cls_id in self.vehicle_classes:
                        # 获取类别名称
                        class_name = self.vehicle_classes[cls_id]

                        # 获取边界框坐标
                        x1, y1, x2, y2 = map(int, box.xyxy[0])

                        # 获取对应类别的颜色
                        color = self.colors.get(class_name, (255, 255, 255))

                        # 绘制边界框
                        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                        # 由于没有可用的中文字体支持，使用 OpenCV 绘制英文标签
                        # 若需要显示中文，仍推荐使用原有的批量添加中文文本方法
                        # 这里简单实现使用 OpenCV 叠加类别和置信度标签
                        # label = f"{class_name}: {confidence:.2f}"
                        # (text_width, text_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)

                        # # 计算文本背景框位置
                        # label_x = x1
                        # label_y = y1 - 5
                        # label_bg_x1 = label_x
                        # label_bg_y1 = label_y - text_height - baseline
                        # label_bg_x2 = label_x + text_width
                        # label_bg_y2 = label_y

                        # # 绘制文本背景
                        # cv2.rectangle(frame, (label_bg_x1, label_bg_y1), (label_bg_x2, label_bg_y2), (0,0,0), -1)

                        # # 绘制文本
                        # cv2.putText(frame, label, (label_x, label_y - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

                        # 准备类别和置信度标签文本信息
                        label = f"{class_name}: {confidence:.2f}"
                        
                        # 将文本信息添加到列表中，留待批量处理
                        texts_info.append((label, (x1, y1 - 5), color, self.font_size, None))

            # 计算并显示FPS（减少FPS显示的频率，避免频繁转换）
            if show_fps:
                self._calculate_fps()
                # 每几帧才更新一次FPS显示
                if self.frame_count % self.fps_display_interval == 0:
                    fps_text = f"FPS: {self.fps:.2f}"
                    # 将FPS文本信息添加到列表中
                    texts_info.append((fps_text, (10, 30), (0, 255, 255), 20, None))

            # 批量添加所有中文文本，只进行一次图像格式转换和一次回转换
            if texts_info:
                  frame = self.batch_add_chinese_text(frame, texts_info)

            # 返回处理后的图像
            return frame

        except Exception as e:
            print(f"检测过程中发生错误: {str(e)}")
            # 出错时返回原始图像的拷贝以避免花屏
            return copy.deepcopy(frame) if isinstance(frame, np.ndarray) else None

    def save_result(self, image, output_path):
        """
        保存处理后的图像
        :param image: 处理后的图像数组
        :param output_path: 输出路径
        :return: 保存成功返回True，否则返回False
        """
        try:
            if image is None or image.size == 0:
                print("错误: 没有可保存的图像")
                return False

            # 创建输出目录（如果不存在）
            output_dir = os.path.dirname(output_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)

            # 保存图像
            cv2.imwrite(output_path, image)
            print(f"结果已保存至: {output_path}")
            return True

        except Exception as e:
            print(f"保存图像时发生错误: {str(e)}")
            return False




# if __name__ == "__main__":
#     # 示例1: 使用文件路径模式
#     detector = VehicleDetector(model_path="yolov8s.pt", device="0")
#     result_image = detector.detect_and_draw("test.jpg", conf_threshold=0.3)
#     if result_image is not None:
#         detector.save_result(result_image, "result.jpg")
#         cv2.imshow("Vehicle Detection (File Mode)", result_image)
#         cv2.waitKey(0)
#         cv2.destroyAllWindows()

    # 示例2: 模拟RTSP帧处理模式（实际使用时会由RTSP解码工具类调用）
    # 这里仅作演示，实际应用中不需要手动调用
    # processor = RTSPProcessorExample()
    # 假设raw_image是从RTSP流获取的原始数据
    # processed_frame = processor.process_rtsp_frame(raw_image)
